from flask import *
from flask_cors import CORS
import torch
import numpy as np
import sys
# 将 YOLOv5 根目录添加到系统路径（根据你的实际路径修改）
sys.path.append(r'D:\course\智能设备开发\期末作业\模型\yolov5-6.0')

from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.augmentations import letterbox
from utils.torch_utils import select_device
from utils.datasets import letterbox
import cv2
from random import randint
from utils.model import  *

app = Flask(__name__)  ##标志该脚本为程序的根目录
# 添加 VideoCamera 类
class VideoCamera:
    def __init__(self, video_path=None):
        if video_path:
            self.cap = cv2.VideoCapture(video_path)
        else:
            self.cap = cv2.VideoCapture(0)  # 摄像头
    
    def get_frame(self):
        success, frame = self.cap.read()
        if success:
            # 这里可以添加 YOLOv5 目标检测逻辑
            return cv2.imencode('.jpg', frame)[1].tobytes()
        return None

# 初始化全局变量
cap = None

# 文件上传处理增加容错
@app.route('/detect', methods=['POST'])
def detect():
    global cap
    if request.method == 'POST':
        if 'videoFile' not in request.files:
            return 'No file uploaded', 400
        file = request.files['videoFile']
        if file.filename == '':
            return 'No selected file', 400
        temp_video_path = 'temp_video.mp4'
        file.save(temp_video_path)  # 更可靠的文件保存方式
        cap = cv2.VideoCapture(temp_video_path)
        return redirect('/')
    return 'Method not allowed', 405

@app.route('/')
def index():
    return render_template('index.html')  #template文件夹下的index.html
def gen(camera):
    while True:
        frame = camera.get_frame()
        # 使用generator函数输出视频流， 每次请求输出的content类型是image/jpeg
        if frame is None:
            break
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')
 
@app.route('/video_feed')  # 这个地址返回视频流响应
def video_feed():
    return Response(gen(VideoCamera()),
                    mimetype='multipart/x-mixed-replace; boundary=frame')
 
@app.route('/detect',methods=['POST'])
def vedio_stream():
    global cap
    if request.method == 'POST':
        stream = request.files['videoFile'].stream
        # 将文件流写入临时文件
        temp_video_path = 'temp_video.mp4'
        with open(temp_video_path, 'wb') as temp_video_file:
            temp_video_file.write(stream.read())
        cap = cv2.VideoCapture(temp_video_path)
    else:
        return 'request method error!'
    return redirect('/')

if __name__ == '__main__':
    app.run(host="127.0.0.1", port=8080, debug=True)  